#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# Licensed under the CANN Open Software License Agreement Version 1.0.
# See LICENSE in the root of the software repository for full license details.

"""
GELU Golden Data Generator

This script generates random input data and corresponding GELU output data (golden reference)
for testing and validation purposes. The data is saved in binary format.
"""

import os
import argparse
import numpy as np
import torch
import torch.nn as nn
from pathlib import Path
from typing import Tuple


def generate_gelu_data(
    input_shape: Tuple[int, ...] = (100,),
    input_range: Tuple[float, float] = (-1.0, 1.0),
    dtype: np.dtype = np.float16,
    output_dir: str = "./data"
) -> None:
    """
    Generate random input and corresponding GELU output data.
    
    Args:
        input_shape: Shape of the input tensor
        input_range: Range of random values for input data
        dtype: Data type for input/output
        output_dir: Directory to save generated data
    """
    # Create output directories
    data_dir = Path(output_dir)
    input_dir = data_dir / "input"
    output_dir = data_dir / "output"
    
    input_dir.mkdir(parents=True, exist_ok=True)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Generate random input data
    np.random.seed(42)  # For reproducibility
    input_x = np.random.uniform(
        low=input_range[0],
        high=input_range[1],
        size=input_shape
    ).astype(dtype)

    # Compute GELU output
    gelu = nn.GELU()
    with torch.no_grad():
        golden = gelu(torch.tensor(input_x)).numpy().astype(dtype)

    # Save data
    input_path = input_dir / "input_x.bin"
    golden_path = output_dir / "golden.bin"
    
    input_x.tofile(input_path)
    golden.tofile(golden_path)
    
    print(f"Input data saved to: {input_path}")
    print(f"Golden output saved to: {input_path}")


def main():
    """Command line interface for data generation."""
    parser = argparse.ArgumentParser(
        description="Generate GELU test data",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument(
        "--shape",
        type=int,
        nargs="+",
        default=[100],
        help="Shape of input tensor"
    )
    parser.add_argument(
        "--range",
        type=float,
        nargs=2,
        default=[-1.0, 1.0],
        metavar=("MIN", "MAX"),
        help="Range of input values"
    )
    parser.add_argument(
        "--dtype",
        choices=["float16", "float32"],
        default="float16",
        help="Data type for input/output"
    )
    parser.add_argument(
        "--output-dir",
        default="./data",
        help="Output directory for generated data"
    )
    
    args = parser.parse_args()
    
    generate_gelu_data(
        input_shape=tuple(args.shape),
        input_range=tuple(args.range),
        dtype=np.dtype(args.dtype),
        output_dir=args.output_dir
    )


if __name__ == "__main__":
    main()